package limit

import (
	"time"
)

// TokenBucket 令牌桶
type TokenBucket struct {
	capacity     int       // 令牌桶容量
	tokens       int       // 当前令牌数量
	refillRate   int       // 令牌填充速率（每秒填充的令牌数）
	lastRefill   time.Time // 上次填充令牌的时间
	OkCallBack   func()    // 令牌桶满时的回调函数
	FailCallBack func()    // 令牌桶空时的回调函数

}

// TokenBucketConfig 令牌桶配置
type TokenBucketConfig struct {
	// 令牌桶容量
	Capacity int
	// 令牌填充速率（每秒填充的令牌数）
	RefillRate int
	// 令牌桶满时的回调函数
	OkCallBack func()
	// 令牌桶空时的回调函数
	FailCallBack func()
}

// NewTokenBucket 创建一个令牌桶
func NewTokenBucket(capacity int, refillRate int) *TokenBucket {
	return &TokenBucket{
		capacity:   capacity,
		tokens:     capacity,
		refillRate: refillRate,
		lastRefill: time.Now(),
	}
}

// TakeToken 从令牌桶中取出一个令牌
func (tb *TokenBucket) TakeToken() {
	now := time.Now()
	// 计算从上次填充令牌到现在应该增加的令牌数量
	elapsed := now.Sub(tb.lastRefill).Seconds()
	tokensToAdd := int(elapsed * float64(tb.refillRate))
	// 更新令牌数量
	tb.tokens = min(tb.capacity, tb.tokens+tokensToAdd)
	tb.lastRefill = now
	// 如果有令牌则扣除一个并返回 true，否则返回 false
	if tb.tokens > 0 {
		tb.tokens--
		tb.OkCallBack()
	} else {
		tb.FailCallBack()
	}
}

// SimpleApiLimit 实现一个简单的接口限流
func SimpleApiLimit(tokenBucketConfig *TokenBucketConfig) *TokenBucket {
	tokenBucket := NewTokenBucket(tokenBucketConfig.Capacity, tokenBucketConfig.RefillRate)
	tokenBucket.FailCallBack = tokenBucketConfig.FailCallBack
	tokenBucket.OkCallBack = tokenBucketConfig.OkCallBack
	tokenBucket.TakeToken()
	return tokenBucket
}
